package com.example.devsa_platform.service.impl;

import co.elastic.clients.elasticsearch._types.aggregations.Aggregation;
import co.elastic.clients.elasticsearch._types.aggregations.Buckets;
import co.elastic.clients.elasticsearch._types.aggregations.StringTermsBucket;
import com.example.devsa_platform.common.Result;
import com.example.devsa_platform.pojo.doc.ArtificialIntelligenceDoc;
import com.example.devsa_platform.pojo.vo.stats.ArtificialIntelligenceVo;
import com.example.devsa_platform.service.ArtificialIntelligenceService;
import jakarta.annotation.Resource;
import org.springframework.data.elasticsearch.client.elc.ElasticsearchAggregations;
import org.springframework.data.elasticsearch.client.elc.ElasticsearchTemplate;
import org.springframework.data.elasticsearch.client.elc.NativeQuery;
import org.springframework.data.elasticsearch.core.AggregationsContainer;
import org.springframework.stereotype.Service;

import java.util.*;
import java.util.stream.Collectors;

@Service
public class ArtificialIntelligenceImpl implements ArtificialIntelligenceService {
    @Resource
    private ElasticsearchTemplate template;

    @Override
    public Result<ArtificialIntelligenceVo> ArtificialIntelligenceStats() {
        NativeQuery query = NativeQuery.builder()
                .withAggregation("aiSelectGroup", Aggregation.of(ag -> ag
                        .terms(te -> te.field("ai_select.keyword")))
                )
                .withAggregation("aiSentGroup", Aggregation.of(ag -> ag
                        .terms(te -> te.field("ai_sent.keyword")))
                )
                .withAggregation("aiAccGroup", Aggregation.of(ag -> ag
                        .terms(te -> te.field("ai_acc.keyword")))
                )
                .withAggregation("aiBenGroup", Aggregation.of(ag -> ag
                        .terms(te -> te.field("ai_ben.keyword")))
                )
                .withAggregation("aiToolInterestInUseGroup", Aggregation.of(ag -> ag
                        .terms(te -> te.field("ai_tool_interest_in_use.keyword")))
                )
                .withAggregation("aiToolCurrentlyUsingGroup", Aggregation.of(ag -> ag
                        .terms(te -> te.field("ai_tool_currently_using.keyword")))
                )
                .withAggregation("aiToolNotInterestedInUsingGroup", Aggregation.of(ag -> ag
                        .terms(te -> te.field("ai_tool_not_interested_in_using.keyword")))
                )
                .withAggregation("aiNextVeryDifferentGroup", Aggregation.of(ag -> ag
                        .terms(te -> te.field("ai_next_very_different.keyword")))
                )
                .withAggregation("aiNextNeitherDifferentNorSimilarGroup", Aggregation.of(ag -> ag
                        .terms(te -> te.field("ai_next_neither_different_nor_similar.keyword")))
                )
                .withAggregation("aiNextSomewhatSimilarGroup", Aggregation.of(ag -> ag
                        .terms(te -> te.field("ai_next_somewhat_similar.keyword")))
                )
                .withAggregation("aiNextVerySimilarGroup", Aggregation.of(ag -> ag
                        .terms(te -> te.field("ai_next_very_similar.keyword")))
                )
                .withAggregation("aiNextSomewhatDifferentGroup", Aggregation.of(ag -> ag
                        .terms(te -> te.field("ai_next_somewhat_different.keyword")))
                )
                .build();

        AggregationsContainer<?> aggregationsContainer = template.search(query, ArtificialIntelligenceDoc.class).getAggregations();
        ArtificialIntelligenceVo artificialIntelligenceVo = null;

        if (aggregationsContainer instanceof ElasticsearchAggregations elasticsearchAggregations) {
            List<Map<String, Object>> aiSelectGroup = parseAndAggregateResults("aiSelectGroup", elasticsearchAggregations);
            List<Map<String, Object>> aiSentGroup = parseAndAggregateResults("aiSentGroup", elasticsearchAggregations);
            List<Map<String, Object>> aiAccGroup = parseAndAggregateResults("aiAccGroup", elasticsearchAggregations);
            List<Map<String, Object>> aiBenGroup = parseAndAggregateResults("aiBenGroup", elasticsearchAggregations);
            List<Map<String, Object>> aiToolInterestInUseGroup = parseAndAggregateResults("aiToolInterestInUseGroup", elasticsearchAggregations);
            List<Map<String, Object>> aiToolCurrentlyUsingGroup = parseAndAggregateResults("aiToolCurrentlyUsingGroup", elasticsearchAggregations);
            List<Map<String, Object>> aiToolNotInterestedInUsingGroup = parseAndAggregateResults("aiToolNotInterestedInUsingGroup", elasticsearchAggregations);
            List<Map<String, Object>> aiNextVeryDifferentGroup = parseAndAggregateResults("aiNextVeryDifferentGroup", elasticsearchAggregations);
            List<Map<String, Object>> aiNextNeitherDifferentNorSimilarGroup = parseAndAggregateResults("aiNextNeitherDifferentNorSimilarGroup", elasticsearchAggregations);
            System.out.println(aiNextNeitherDifferentNorSimilarGroup);
            List<Map<String, Object>> aiNextSomewhatSimilarGroup = parseAndAggregateResults("aiNextSomewhatSimilarGroup", elasticsearchAggregations);
            List<Map<String, Object>> aiNextVerySimilarGroup = parseAndAggregateResults("aiNextVerySimilarGroup", elasticsearchAggregations);
            List<Map<String, Object>> aiNextSomewhatDifferentGroup = parseAndAggregateResults("aiNextSomewhatDifferentGroup", elasticsearchAggregations);

            artificialIntelligenceVo = new ArtificialIntelligenceVo(
                    aiSelectGroup,
                    aiSentGroup,
                    aiAccGroup,
                    aiBenGroup,
                    aiToolInterestInUseGroup,
                    aiToolCurrentlyUsingGroup,
                    aiToolNotInterestedInUsingGroup,
                    aiNextVeryDifferentGroup,
                    aiNextNeitherDifferentNorSimilarGroup,
                    aiNextSomewhatSimilarGroup,
                    aiNextVerySimilarGroup,
                    aiNextSomewhatDifferentGroup
            );

        }
        System.out.println(artificialIntelligenceVo);
        return Result.success(artificialIntelligenceVo);
    }

    private List<Map<String, Object>> parseAndAggregateResults(String aggregationName, ElasticsearchAggregations elasticsearchAggregations) {
        Map<String, Long> aggregatedCounts = new HashMap<>();

        List<Map<String, Object>> originalResults = parseAggregation(aggregationName, elasticsearchAggregations);
        for (Map<String, Object> result : originalResults) {
            String key = (String) result.get("key");
            Long docCount = ((Number) result.get("doc_count")).longValue();

            // Split the key by ";"
            String[] parts = key.split(";");
            for (String part : parts) {
                part = part.trim();
                aggregatedCounts.put(part, aggregatedCounts.getOrDefault(part, 0L) + docCount);
            }
        }

        // Convert aggregatedCounts to the desired output
        List<Map<String, Object>> parsedResults = new ArrayList<>();
        for (Map.Entry<String, Long> entry : aggregatedCounts.entrySet()) {
            Map<String, Object> map = new HashMap<>();
            map.put("key", entry.getKey());
            map.put("doc_count", entry.getValue());
            parsedResults.add(map);
        }

        return parsedResults;
    }

    private List<Map<String, Object>> parseAggregation(String aggName, ElasticsearchAggregations elasticsearchAggregations) {
        Buckets<StringTermsBucket> buckets = Objects.requireNonNull(elasticsearchAggregations.get(aggName)).aggregation().getAggregate().sterms().buckets();
        return buckets.array().stream()
                .map(bucket -> {
                    String key = bucket.key().stringValue();
                    long docCount = bucket.docCount();
                    Map<String, Object> map = new HashMap<>();
                    map.put("key", key);
                    map.put("doc_count", docCount);
                    return map;
                })
                .collect(Collectors.toList());
    }

    }
